import pybullet as p
import pybullet_data
import numpy as np


p.connect(p.GUI)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

p.resetSimulation(p.RESET_USE_DEFORMABLE_WORLD)
p.setGravity(0, 0, -9.8)
    
    
ur4_urdf = 'ur5e/ur5e.urdf'


p.loadURDF("plane.urdf", [0, 0, -0.001])

# robot_id = p.loadURDF(ur4_urdf, basePosition=[0,0,0.1], useFixedBase=True)
robot_id = p.loadURDF("ur5e/ur5e.urdf", [0, 0, 0], flags=p.URDF_USE_MATERIAL_COLORS_FROM_MTL)
num_joints = p.getNumJoints(robot_id)
print(f"⚙️ 总关节数: {num_joints}")


joint_ids = [p.getJointInfo(robot_id, i) for i in range(p.getNumJoints(robot_id))]
joint_ids = [j[0] for j in joint_ids if j[2] == p.JOINT_REVOLUTE]

for joint_idx in range(num_joints):
    # 基础信息获取
    info = p.getJointInfo(robot_id, joint_idx)
    if info[2] == 4:
        continue
    print(f'Joint index {joint_idx}, {info[1].decode("utf-8")}, {info[2]}, {info[7]}')

home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0)
home_ee_euler = (np.pi, 0, np.pi)

for i in range(len(joint_ids)):
    p.resetJointState(robot_id, joint_ids[i], home_joints[i])


while  True:
    p.stepSimulation()
    pass

